library(tidyr)
library(dplyr)
library(ggplot2)
library(ggpubr)

SBS_earlylate <- read.csv("ExtendedDataFigure_4_data.csv")

# Calculate differences between early and late
SBS_early <- SBS_earlylate %>% filter(time == "early")
SBS_late <- SBS_earlylate %>% filter(time == "late")

#### Plot activities ####

SBS_dif <- SBS_early[, -c(1, 2)] - SBS_late[, -c(1, 2)]
SBS_dif <- SBS_dif %>%
  mutate_all(~ case_when(. > 0.06 ~ "early", . < -0.06 ~ "late", TRUE ~ "unclear")) %>%
  mutate_all(~ factor(., levels = c("unclear", "early", "late")))

# Set row and column names
rownames(SBS_dif) <- SBS_early$donor_id
colnames(SBS_dif) <- paste0(colnames(SBS_dif), "_time")
SBS_dif <- SBS_dif %>% tibble::rownames_to_column("donor_id")

# Merge differences with original data
SBS_plot <- SBS_earlylate %>% left_join(SBS_dif, by = "donor_id")

# Function to plot timing data
plot_timing <- function(SBS_plot, signatures, title = NA, print = TRUE, save = TRUE) {
  # Perform Wilcoxon paired test for each signature
  output_stats <- NULL
  for (sig in signatures) {
    wilcox <- wilcox.test(get(sig) ~ time, data = SBS_plot, paired = T)
    stats <- data.frame(signature = sig, p_val = wilcox$p.value)
    output_stats <- rbind(output_stats, stats)
  }

  output_stats <- output_stats %>% mutate(p_adj = p.adjust(p_val, method = "BH"))

  for (sig in signatures) {
    # filter cases positive for early and/or late activities
    SBS_plot.pos_bysample <- SBS_plot %>%
      select("donor_id", "time", sig) %>%
      spread(time, sig) %>%
      filter(!(early == 0 & late == 0)) 
    
    SBS_plot.pos <- SBS_plot.pos_bysample %>%
      gather(time, sig, -donor_id)
    
    qval <- output_stats %>%
      filter(signature == sig) %>%
      pull(p_adj)

    plot_title <- ifelse(is.na(title),
      paste0(sig, " (", nrow(SBS_plot.pos_bysample), "/", nrow(SBS_plot) / 2, ")"),
      paste0(sig, " ", title, " (", nrow(SBS_plot.pos_bysample), "/", nrow(SBS_plot) / 2, ")")
    )

    p <- SBS_plot %>%
      arrange(get(paste0(sig, "_time"))) %>%
      mutate(donor_id = factor(donor_id, levels = unique(donor_id))) %>%
      ggplot(aes(
        x = time, y = get(sig)
      )) +
      geom_line(aes(group = donor_id, color = factor(get(paste0(sig, "_time")))), alpha = 0.6) +
      geom_boxplot(
        data = SBS_plot.pos,
        aes(x = time, y = sig, fill = time),
        width = .2, color = "#404040",
        alpha = 0.25, linewidth = .2, staplewidth = .2, outlier.shape = NA
      ) +
      geom_point(aes(group = donor_id), size = .5, alpha = 0.25) +
      scale_x_discrete(expand = c(.1, .1)) +
      scale_y_continuous(limits = c(0, 1)) +
      scale_color_manual(values = c("early" = "#6495ED", "late" = "#E97132", "unclear" = "#7D7D7D")) +
      scale_fill_manual(values = c("early" = "#6495ED", "late" = "#E97132")) +
      labs(
        title = plot_title,
        subtitle = paste0("q-val = ", signif(qval, 4)),
        x = "", y = "relative activity", color = ""
      ) +
      theme_minimal() +
      theme(
        legend.position = "none",
        plot.title = element_text(face = "bold", size = 12, hjust = 0.5),
        plot.subtitle = element_text(size = 10, hjust = 0.5),
        axis.title = element_text(size = 12),
        axis.text = element_text(size = 12),
        axis.line = element_line(size = .3,color = "#404040"),
        axis.ticks = element_line(size = .3,color = "#404040"),
        panel.grid = element_blank()
      )

    if (print == TRUE) {
      print(p)
    }

    if (save == TRUE) {
      if (is.na(title)) {
        ggsave(paste0("output/ExtendedDataFigure_4_", sig, "_", Sys.Date(), ".pdf"), device = "pdf", width = 3.25, height = 3.5, dpi = 700)
      } else {
        ggsave(paste0("output/ExtendedDataFigure_4_", sig, "_", title, "_", Sys.Date(), ".pdf"), device = "pdf", width = 3.25, height = 3.5, dpi = 700)
      }
    }
  }
}

# Plot for the whole dataset
signatures <- names(SBS_earlylate[-c(1, 2)])
plot_timing(SBS_plot, signatures, print = T, save = T)

# Load metadata for sub-group analysis
source("HNC_metadata_tidy.R")

data <- data %>%
  mutate(
    subsite = factor(subsite, levels = c("Larynx", "OC", "OPC", "Hypopharynx")),
    tobacco_ever = factor(tobacco_ever, labels = c("Never-smoker", "Ever-smoker")),
    argentina = ifelse(country == "Argentina", "yes", "no"),
    greece = ifelse(country == "Greece", "yes", "no"),
    italy = ifelse(country == "Italy", "yes", "no"),
    brazil = ifelse(country == "Brazil", "yes", "no"),
    czechRepublic = ifelse(country == "Czech Republic", "yes", "no"),
    romania = ifelse(country == "Romania", "yes", "no"),
    age_group = factor(age_group, levels = c("55-65", "0-45", "45-55", "65-75", "75+")),
    age = scale(age_diag)
  )

# Plot for each level of selected variables
vars <- c("tobacco_ever","subsite")
signatures  <- "SBS_I"

for (v in vars) {
  for (lv in levels(data[, v])) {
    data.lv <- data[data[, v] == lv, ]
    SBS_plot.lv <- SBS_plot[SBS_plot$donor_id %in% data.lv$donor_id, ]

    plot_timing(SBS_plot.lv, signatures, title = lv, print = T, save = T)
  }
}

#### Logistic regression ####

# Load required script for data handling functions
source("data_handling.R")
source("risk_factor_regressions.R")

# Merge early and late activities with metadata and dichotomize signatures
SBS.early <- SBS_early %>% select(-time)
SBS.late <- SBS_late %>% select(-time)

data.early <- data %>% right_join(SBS.early, by = "donor_id")
data.early <- dichotomizeSigs(metadata = data.early, sigsn = SBS.early)

data.late <- data %>% right_join(SBS.late, by = "donor_id")
data.late <- dichotomizeSigs(metadata = data.late, sigsn = SBS.late)

# Define regression parameters and confounders
parameters <- c(
  "tobacco_ever",
  "OC", "OPC", "Larynx", "Hypopharynx",
  "argentina", "greece", "italy", "brazil", "czechRepublic", "romania"
)

confounders <- c("age", "sex", "region", "subsite", "alcohol_ever", "tobacco_ever")

# Run logistic regression models
regression.early <- LRmodel(data.early,
  signatures = c("SBS_I_cat", "SBS4_cat", "SBS92_cat"),
  parameters, confounders, pval = 1, make_plot = F
)

regression.late <- LRmodel(data.late,
  signatures = c("SBS_I_cat", "SBS4_cat", "SBS92_cat"),
  parameters, confounders, pval = 1, make_plot = F
)

# Combine regression results and prepare for plotting
bubbleplot <- bind_rows(
  regression.early$regression_param %>% mutate(signature = str_replace(signature, "_cat", " early")),
  regression.late$regression_param %>% mutate(signature = str_replace(signature, "_cat", " late"))
) %>%
  mutate(
    signature = factor(signature),
    independent_vars = factor(str_remove(independent_vars, "yes|tobacco_ever")),
    log_OR = log2(OR),
    # Modifying p_adj so that it's corrected for comparisons across all signatures
    p_adj = ifelse(p_val * length(signatures) < 1, p_val * length(signatures), 1),
    p_adj_log = -log10(p_adj),
    signif = ifelse(p_adj < 0.05, "p adj < 0.05", "p adj > 0.05"),
    signif = factor(signif)
  )

# Create bubble plot for tobacco and alcohol associations
p_tob <- bubbleplot %>%
  filter(independent_vars == "Ever-smoker") %>%
  mutate(independent_vars = factor(independent_vars, labels = "Tobacco")) %>%
  ggplot(aes(
    y = independent_vars, x = signature,
    color = log_OR, size = p_adj_log, shape = signif
  )) +
  geom_point(stroke = 0.75) +
  scale_color_gradient2(
    high = "#ca0020", mid = "white", low = "#0571b0", midpoint = 0, limits = c(min(bubbleplot$log_OR), max(bubbleplot$log_OR)) # limits = c(-4,4),breaks = c(0,2,4,6)
  ) +
  scale_shape_manual(values = c("p adj < 0.05" = 19, "p adj > 0.05" = 21)) +
  scale_size_continuous(limits = c(0, max(bubbleplot$p_adj_log))) +
  labs(
    title = "Associations with epidemiological variables",
    subtitle = "Corrected for sex, age of diagnosis, subsite, region, \nand alcohol consumption",
    x = "", y = "", color = "log2(OR)", size = "-log10(p adj)", shape = ""
  ) +
  theme_bw() +
  theme(
    legend.position = "right",
    plot.title = element_text(face = "bold", size = 12),
    plot.subtitle = element_text(size = 10),
    axis.text = element_text(size = 12),
    axis.text.x = element_blank(),
    legend.text = element_text(size = 12),
    panel.grid = element_blank()
  )

p_site <- bubbleplot %>%
  filter(independent_vars %in% c("OC", "OPC", "Larynx", "Hypopharynx")) %>%
  mutate(independent_vars = factor(independent_vars, levels = rev(c("OC", "OPC", "Larynx", "Hypopharynx")))) %>%
  ggplot(aes(
    y = independent_vars, x = signature,
    color = log_OR, size = p_adj_log, shape = signif
  )) +
  geom_point(stroke = 0.75) +
  scale_color_gradient2(
    high = "#ca0020", mid = "white", low = "#0571b0", midpoint = 0, limits = c(min(bubbleplot$log_OR), max(bubbleplot$log_OR)) # limits = c(-4,4),breaks = c(0,2,4,6)
  ) +
  scale_shape_manual(values = c("p adj < 0.05" = 19, "p adj > 0.05" = 21)) +
  scale_size_continuous(limits = c(0, max(bubbleplot$p_adj_log))) +
  labs(
    title = "Associations with epidemiological variables",
    subtitle = "Corrected for sex, age of diagnosis, region, tobacco, \nand alcohol consumption",
    x = "", y = "", color = "log2(OR)", size = "-log10(p adj)", shape = ""
  ) +
  theme_bw() +
  theme(
    legend.position = "right",
    plot.title = element_blank(),
    plot.subtitle = element_text(size = 10),
    axis.text = element_text(size = 12),
    axis.text.x = element_text(angle = 45, hjust = 1),
    legend.text = element_text(size = 12),
    panel.grid = element_blank()
  )

# Combine plots into a single figure and save
p_all <- ggarrange(p_tob, p_site,
  common.legend = TRUE, legend = "right",
  ncol = 1, align = "v", heights = c(0.35, 0.65)
)

ggsave(last_plot(),
  filename = paste0("output/ExtendedDataFigure_4_", Sys.Date(), ".pdf"),
  device = "pdf",
  width = 5.5, height = 4, dpi = 700
)

